"""
向量存储构建器
负责文档加载、处理、向量化并存储到 Qdrant 向量数据库
"""

import os
import logging
from typing import List, Optional
from pathlib import Path

from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    Settings,
    Document
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.embeddings.dashscope import DashScopeEmbedding
from llama_index.llms.dashscope import DashScope
from qdrant_client import QdrantClient
from qdrant_client.http import models
from dotenv import load_dotenv

from ..models.llm_config import LLMConfig
from ..utils.logger import setup_logger

logger = setup_logger(__name__)


# 加载环境变量
load_dotenv()

class VectorStoreBuilder:
    """向量存储构建器"""
    
    def __init__(self, config: LLMConfig):
        """
        初始化向量存储构建器
        
        Args:
            config: LLM 配置对象
        """
        self.config = config
        self.qdrant_client = None
        self.vector_store = None
        self._setup_llm_settings()
        self._setup_qdrant()
    
    def _setup_llm_settings(self):
        """设置 LLM 和嵌入模型"""
        try:
            # 设置嵌入模型
            embed_model = DashScopeEmbedding(
                model_name=os.getenv("EMBEDDING_MODEL"),
                api_key=self.config.api_key
            )
            
            # 设置 LLM
            llm = DashScope(
                model=os.getenv("LLM_MODEL"),
                api_key=self.config.api_key
            )
            
            # 配置全局设置
            Settings.llm = llm
            Settings.embed_model = embed_model
            
            logger.info("LLM 和嵌入模型配置成功")
            
        except Exception as e:
            logger.error(f"配置 LLM 和嵌入模型失败: {e}")
            raise
    
    def _setup_qdrant(self):
        """设置 Qdrant 连接"""
        try:
            # 连接 Qdrant
            self.qdrant_client = QdrantClient(
                host=self.config.qdrant_host,
                port=self.config.qdrant_port
            )
            
            # 创建向量存储
            self.vector_store = QdrantVectorStore(
                client=self.qdrant_client,
                collection_name=self.config.collection_name
            )
            
            logger.info(f"Qdrant 连接成功: {self.config.qdrant_host}:{self.config.qdrant_port}")
            
        except Exception as e:
            logger.error(f"连接 Qdrant 失败: {e}")
            raise
    
    def _clear_collection(self):
        """清空现有集合"""
        try:
            # 检查集合是否存在
            collections = self.qdrant_client.get_collections()
            collection_names = [col.name for col in collections.collections]
            
            if self.config.collection_name in collection_names:
                logger.info(f"清空现有集合: {self.config.collection_name}")
                self.qdrant_client.delete_collection(self.config.collection_name)
            
        except Exception as e:
            logger.warning(f"清空集合时出现警告: {e}")
    
    def _load_documents(self, data_dir: str) -> List[Document]:
        """
        加载文档
        
        Args:
            data_dir: 数据目录路径
            
        Returns:
            文档列表
        """
        try:
            data_path = Path(data_dir)
            if not data_path.exists():
                raise FileNotFoundError(f"数据目录不存在: {data_dir}")
            
            # 使用 SimpleDirectoryReader 加载文档
            reader = SimpleDirectoryReader(
                input_dir=str(data_path),
                required_exts=[".pdf"],
                recursive=True
            )
            
            documents = reader.load_data()
            logger.info(f"成功加载 {len(documents)} 个文档")
            
            return documents
            
        except Exception as e:
            logger.error(f"加载文档失败: {e}")
            raise
    
    def _process_documents(self, documents: List[Document]) -> List[Document]:
        """
        处理文档（切分等）
        
        Args:
            documents: 原始文档列表
            
        Returns:
            处理后的文档列表
        """
        try:
            # 创建文档切分器
            text_splitter = SentenceSplitter(
                chunk_size=self.config.chunk_size,
                chunk_overlap=self.config.chunk_overlap
            )
            
            # 处理文档
            processed_docs = []
            for doc in documents:
                # 切分文档
                nodes = text_splitter.get_nodes_from_documents([doc])
                for node in nodes:
                    processed_docs.append(Document(
                        text=node.text,
                        metadata=node.metadata
                    ))
            
            logger.info(f"文档处理完成，生成 {len(processed_docs)} 个文档片段")
            return processed_docs
            
        except Exception as e:
            logger.error(f"处理文档失败: {e}")
            raise
    
    def build_vector_store(self, data_dir: str, force_rebuild: bool = True):
        """
        构建向量存储
        
        Args:
            data_dir: 数据目录路径
            force_rebuild: 是否强制重建（清空现有数据）
        """
        try:
            logger.info("开始构建向量存储...")
            
            # 如果需要重建，先清空集合
            if force_rebuild:
                self._clear_collection()
            
            # 加载文档
            documents = self._load_documents(data_dir)
            if not documents:
                logger.warning("未找到任何文档")
                return
            
            # 处理文档
            processed_docs = self._process_documents(documents)
            
            # 创建存储上下文
            storage_context = StorageContext.from_defaults(
                vector_store=self.vector_store
            )
            
            # 创建向量索引
            logger.info("开始创建向量索引...")
            index = VectorStoreIndex.from_documents(
                processed_docs,
                storage_context=storage_context,
                show_progress=True
            )
            
            logger.info("向量存储构建完成！")
            
        except Exception as e:
            logger.error(f"构建向量存储失败: {e}")
            raise
    
    def get_index(self) -> VectorStoreIndex:
        """
        获取向量索引
        
        Returns:
            向量索引对象
        """
        try:
            storage_context = StorageContext.from_defaults(
                vector_store=self.vector_store
            )
            
            index = VectorStoreIndex.from_vector_store(
                vector_store=self.vector_store,
                storage_context=storage_context
            )
            
            return index
            
        except Exception as e:
            logger.error(f"获取向量索引失败: {e}")
            raise


def main():
    """主函数 - 可直接运行此文件来构建向量存储"""
    try:
        # 加载配置
        config = LLMConfig()
        
        # 创建构建器
        builder = VectorStoreBuilder(config)
        
        # 构建向量存储
        data_dir = os.path.join(os.path.dirname(__file__), "../../data")
        builder.build_vector_store(data_dir, force_rebuild=True)
        
        logger.info("向量存储构建完成！")
        
    except Exception as e:
        logger.error(f"构建向量存储时发生错误: {e}")
        raise


if __name__ == "__main__":
    main()
